# Natural Language Toolkit: Aligned Sentences
#
# Copyright (C) 2001-2010 NLTK Project
# URL: <http://www.nltk.org/>
# For license information, see LICENSE.TXT

import nltk.metrics.scores
from alignment import Alignment, _giza2pair
try:
    import pylab
    __pylab_loaded__ = True
except ImportError:
    __pylab_loaded__ = False


class AlignedSent(object):
    """
    Aligned sentence object.  Encapsulates two sentences along with
    an C{Alignment} between them.
    """
    
    def __init__(self, words = [], mots = [], alignment = '', \
                 possible_alignment = '', encoding = 'latin-1'):
        """
        Initialize a new C{AlignedSent}.
        
        @param words: source language words
        @type words: C{list} of C{str}
        @param mots: target language words
        @type mots: C{list} of C{str}
        @param alignment: the word-level alignments between the source
            and target language
        @type alignment: C{Alignment}
        """
        if not isinstance(alignment, Alignment):
            alignment = Alignment(alignment)
        if not isinstance(possible_alignment, Alignment):
            possible_alignment = Alignment(possible_alignment)
        self._words = words
        self._mots = mots
        self._alignment = alignment
        self._possible_alignment = possible_alignment

    @property
    def words(self):
        return self._words

    @property
    def mots(self):
        return self._mots

    @property
    def alignment(self):
        return self._alignment

    @alignment.setter
    def alignment(self, alignment):
        if not isinstance(alignment, Alignment):
            alignment = Alignment(alignment)
        self._alignment = alignment

    @property
    def possible_alignment(self):
        return self._possible_alignment

    @possible_alignment.setter
    def possible_alignment(self, alignment):
        if not isinstance(alignment, Alignment):
            alignment = Alignment(alignment)
        self._possible_alignment = alignment


    def _check_align(self, a):
        # check legality of alignments
        """
        @param a: alignment to be checked
        @raise: IndexError if alignment is out of sentence boundary
        @return: True if passed alignment check
        @rtype: boolean
        """
        ## curently disabled 
        ##for (i,j) in a:
        ##    if i > len(self._words) - 1 or j > len(self._mots) - 1:
        ##        raise IndexError("Alignment out of sentence boundary")
        return True

        
    def __repr__(self):
        """
        @return: A string representation for this C{AlignedSent}.
        @rtype: C{string}
        """
        return "AlignedSent(%r, %r, %r)" % (self._words, self._mots, self._alignment)

    def gui_repr(self):
        """
        Generate a GUI to represent the sentence alignments
        """
        if __pylab_loaded__:
            fig_width = max(len(self.words), len(self.mots)) + 1
            fig_height = 3
            pylab.figure(figsize=(fig_width*0.8, fig_height*0.8), facecolor='w')
            pylab.box(on=False)
            pylab.subplots_adjust(left=0, right=1, bottom=0, top=1)
            pylab.xlim(-1, fig_width - 1)
            pylab.ylim(0, fig_height)
            pylab.xticks([])
            pylab.yticks([])

            e = [0 for _ in xrange(len(self.words))]
            f = [0 for _ in xrange(len(self.mots))]
            for (i, j) in self.alignment:
                e[i] = 1
                f[j] = 1
                # draw the middle line
                pylab.arrow(i, 2, j - i, -1, color='r')
            for i in xrange(len(e)):
                # draw e side line
                pylab.text(i, 2.5, self.words[i], ha = 'center', va = 'center',
                        rotation=30)
                if e[i] == 1:
                    pylab.arrow(i, 2.5, 0, -0.5, color='r', alpha=0.3, lw=2)
            for i in xrange(len(f)):
                # draw f side line
                 pylab.text(i, 0.5, self.mots[i], ha = 'center', va = 'center',
                        rotation=30)
                 if f[i] == 1:
                    pylab.arrow(i, 0.5, 0, 0.5, color='r', alpha=0.3, lw=2)

            pylab.draw()
    
    @staticmethod
    def block_and_view_gui_repr():
        """
        Block the program to view previous spawned Repr GUIs
        """
        if __pylab_loaded__:
            pylab.show()
    
    def __str__(self):
        """
        @return: A string representation for this C{AlignedSent}.
        @rtype: C{string}
        """
        source = " ".join(self._words)[:20] + "..."
        target = " ".join(self._mots)[:20] + "..."
        return "<AlignedSent: '%s' -> '%s'>" % (source, target)

    def update_align(self, e,f, s = True):
        """to add a new alignment.
        @param e: index of the alignment to be updated.
        @type e: int
        @param f: index of the alignment to be udpated.
        @type f: int
        """
        a = (e,f)
        self._check_align(Alignment([a]))
        if s:
            self._alignment = self._alignment.union(Alignment([a]))
            self._alignment = self._alignment.union(Alignment([a]))
        else:
            self.remove_align(e,f)
        self._possible_alignment = self._possible_alignment.union(Alignment([a]))

    def remove_align(self, e,f):
        """remove target Alignment with key (e,f).
        """
        try:
            self._alignment = self._alignment.difference([(e,f)])
            self._possible_alignment = \
                        self._possible_alignment.difference([(e,f)])
        except KeyError:
            return

    def invert(self):
        """ 
        @return: the invert object
        @rtype: AlignedSent
        """
        if self._possible_alignment == None:
            return AlignedSent(self._mots, self._words, \
                               self._alignment.invert())

        return AlignedSent(self._mots, self._words, \
                           self._alignment.invert(), \
                           self._possible_alignment.invert())
    
def precision(aligned_sent, reference_aligned_sent):
    """Calculates the precision of an aligned sentence with respect to a "gold 
    standard" reference alignment.

    The "possible" precision is used since it doesn't penalise for finding an 
    alignment that was marked as "possible".
    """

    # Get alignments in set of 2-tuples form
    a = set(((s[0], s[1]) for s in aligned_sent.alignment))
    p = set(((s[0], s[1]) for s in reference_aligned_sent.possible_alignment))

    # Call NLTKs existing functions for precision
    return nltk.metrics.scores.precision(p, a)


def recall(aligned_sent, reference_aligned_sent):
    """Calculates the recall of an aligned sentence with respect to a "gold 
    standard" reference alignment.

    The "sure" recall is used so we don't penalise for missing an alignment 
    that was only marked as "possible".
    """

    # Get alignments in set of 2-tuples form
    a = set(((s[0], s[1]) for s in aligned_sent.alignment))
    s = set(((s[0], s[1]) for s in reference_aligned_sent.alignment))

    # Call NLTKs existing functions for recall
    return nltk.metrics.scores.recall(sure_align, align)


def alignment_error_rate(aligned_sent, reference_aligned_sent):
    """Calculates the Alignment Error Rate (AER) of an aligned sentence with 
    respect to a "gold standard" reference alignment.

    Return an error rate between 0.0 (perfect alignment) and 1.0 (no 
    alignment).
    """

    # Get alignments in set of 2-tuples form
    a = set(((s[0], s[1]) for s in aligned_sent.alignment))
    s = set(((s[0], s[1]) for s in reference_aligned_sent.alignment))
    p = set(((s[0], s[1]) for s in reference_aligned_sent.possible_alignment))

    # Sanity check
    assert(s.issubset(p))

    # Return the Alignment Error Rate
    return 1.0 - float(len(a & s) + len(a & p)) / float(len(a) + len(s))

